from const import AudioDatasetDir_Original, ProjectDir
from dataset_utils import Audio
import dataset_utils

import os
import unittest
import cv2

# 将 Audio generate 出来的图像保存到指定位置，之后可以用肉眼检查正确性
IMG_SAVE_FLAG = True

class Unittest_for_Audio(unittest.TestCase):
    
    def setUp(self) -> None:
        super().setUp()
        self.test_filenames = set()
        
        files = os.listdir(AudioDatasetDir_Original)
        for f in files:
            f_info = f.strip().split('.')[0].split('_')
            # f_info = ['226', '1b1', 'Pl', 'sc', 'LittC2SE'], f='226_1b1_Pl_sc_LittC2SE.wav'
            if len(f_info) != 5:
                # 该目录下还有一些其他的文件，直接忽视就好
                continue
            
            self.test_filenames.add('_'.join(f_info))

        # print(self.test_filenames)
        
        self.SAVE_DIR = os.path.join(ProjectDir, 'unittest')
        if not os.path.exists(self.SAVE_DIR):
            os.makedirs(self.SAVE_DIR)
    
    
    def test_generateIMG(self):
        save_flag = dataset_utils.save_detect_img
        dataset_utils.save_detect_img = False
        for fn in self.test_filenames:
            AudioInstance = Audio(fn)
            img, bboxes = AudioInstance.GeneratePic()
            # region check_bboxes
            last_bbox = {}
            for bbox in bboxes:
                if last_bbox != {}:
                    self.assertEqual(bbox['cycle_id'], last_bbox['cycle_id']+1)
                last_bbox = bbox
            last_bbox = {}
            for bbox in bboxes:
                if last_bbox != {}:
                    self.assertEqual(bbox['start_index'], last_bbox['end_index'])
                last_bbox = bbox
            last_bbox = {}
            for bbox in bboxes:
                if last_bbox != {}:
                    self.assertEqual(bbox['start_x'], last_bbox['end_x'])
                last_bbox = bbox
            last_bbox = {}
            for bbox in bboxes:
                if last_bbox != {}:
                    self.assertEqual(bbox['start_y'], last_bbox['start_y'])
                last_bbox = bbox
            last_bbox = {}
            for bbox in bboxes:
                if last_bbox != {}:
                    self.assertEqual(bbox['end_y'], last_bbox['end_y'])
                last_bbox = bbox
            # endregion
            
            if IMG_SAVE_FLAG:
                # 不带框的先画一遍
                cv2.imwrite(f"{self.SAVE_DIR}/NoBBOXes-{AudioInstance.FileName}.jpg", img)
                # 带框的
                for bbox in bboxes:
                    cv2.rectangle(
                        img,
                        (int(bbox['start_x']), int(bbox['start_y'])),
                        (int(bbox['end_x']), int(bbox['end_y'])),
                        (255, 255, 255),
                        1
                    )
                cv2.imwrite(f"{self.SAVE_DIR}/BBOXes-{AudioInstance.FileName}.jpg", img)
                    
        

        dataset_utils.save_detect_img = save_flag            


def Functionality_Test_for_Audio_IMG_Generation():
    pass
    


if __name__ == '__main__':
    unittest.main()
